import gymnasium as gym
import torch
import numpy as np
from dqn_agent import DQNAgent

def load_model(env, model_path):
    """加载训练好的模型"""
    checkpoint = torch.load(model_path)
    agent = DQNAgent(env)
    
    agent.policy_net.load_state_dict(checkpoint['policy_net'])
    agent.target_net.load_state_dict(checkpoint['target_net'])
    agent.optimizer.load_state_dict(checkpoint['optimizer'])
    agent.epsilon = checkpoint['epsilon']
    
    return agent

def long_run_test(model_path, total_steps=100000):
    # 创建修改后的环境
    env = gym.make('CartPole-v1')
    env._max_episode_steps = float('inf')  # 移除默认的200步限制
    
    # 加载智能体
    agent = load_model(env, model_path)
    
    # 初始化跟踪变量
    state, info = env.reset()
    total_reward = 0
    episode_count = 0
    steps_remaining = total_steps
    
    try:
        while steps_remaining > 0:
            # 禁用探索
            action = agent.select_action(state, training=False)
            
            next_state, reward, done, _ , _= env.step(action)
            
            # 更新统计
            total_reward += reward
            steps_remaining -= 1
            
            # 环境渲染（前1000步渲染）
            if total_steps - steps_remaining < 1000:
                env.render()
            
            # 处理环境终止
            if done or steps_remaining <= 0:
                print(f"Episode {episode_count+1} ended after {total_steps - steps_remaining} steps")
                state = env.reset()
                episode_count += 1
            else:
                state = next_state
                
            # 定期报告
            if (total_steps - steps_remaining) % 1000 == 0:
                print(f"Steps: {total_steps - steps_remaining}, Total Reward: {total_reward}")
                
    finally:
        env.close()
        print("\nFinal Statistics:")
        print(f"Total Steps Completed: {total_steps - steps_remaining}")
        print(f"Total Reward Accumulated: {total_reward}")
        print(f"Number of Episodes: {episode_count}")

if __name__ == "__main__":
    model_path = "saved_models/20250127_211755/model_ep199.pth"  # 替换为实际路径
    long_run_test(model_path)